from torch import nn

from NN.nn_module import NNM

class DNN(NNM):

    def __init__(self, model_note, in_channels, output_shape, device_):
        super(DNN, self).__init__(model_note, in_channels, output_shape, device_)
    
    def module(self):
        self.nn_stack = nn.Sequential(
            nn.Conv2d(in_channels=self.in_channels, out_channels=32, kernel_size=3, stride=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, 1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Dropout2d(0.25),
            nn.Flatten(),
            nn.Linear(12544, 128),
            nn.ReLU(),
            nn.Dropout2d(0.5),
            nn.Linear(128, self.y_shape),
            nn.LogSoftmax(dim=1)
        )